{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fd2bc687",
   "metadata": {},
   "source": [
    "# 初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "560725a2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-26T14:55:19.905185800Z",
     "start_time": "2023-05-26T14:55:17.404362900Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/xiaobo/Project/Trans-WSSS\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xiaobo/anaconda3/envs/p10-timm/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "%cd '/home/xiaobo/Project/Trans-WSSS'\n",
    "%pwd\n",
    "\n",
    "import sys\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sys.path.append(\"others/segment_anything\")\n",
    "from segment_anything import sam_model_registry\n",
    "\n",
    "sys.path = ['.', './src'] + sys.path\n",
    "from libs.sam import SamAuto\n",
    "from libs.viz.viz_anns import show_imgs_anns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27c41445",
   "metadata": {},
   "source": [
    "# 读取图片和SAM。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ad354922",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-26T14:55:19.992366200Z",
     "start_time": "2023-05-26T14:55:19.911190500Z"
    }
   },
   "outputs": [],
   "source": [
    "img_ids = ['2007_000032', '2007_000170', '2007_003205',\n",
    "           '2007_009724', '2008_001784', '2008_006844',\n",
    "           '2008_007281', '2010_004844', '2009_004050']\n",
    "\n",
    "imgs = [cv2.cvtColor(cv2.imread(f'datasets/VOC2012/JPEGImages/{img_id}.jpg'), cv2.COLOR_BGR2RGB) for img_id in img_ids]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1848a108",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-26T14:55:39.336638900Z",
     "start_time": "2023-05-26T14:55:24.442987400Z"
    }
   },
   "outputs": [],
   "source": [
    "sam_checkpoint = \"pretrains/SAM/sam_vit_h_4b8939.pth\"\n",
    "model_type = \"vit_h\"\n",
    "\n",
    "device = \"cuda:7\"\n",
    "\n",
    "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "sam = sam.to(device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 建立SAM自动分割器，并推理。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "68364513",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-26T14:57:23.799397200Z",
     "start_time": "2023-05-26T14:57:23.768144300Z"
    }
   },
   "outputs": [],
   "source": [
    "# * 官方Jupyter V2\n",
    "# mask_generator = SamAutomaticMaskGenerator(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "# )\n",
    "# * 默认值.\n",
    "# mask_generator = SamAutomaticMaskGenerator(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.88,\n",
    "#     stability_score_thresh=0.95,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=0,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=1,\n",
    "#     min_mask_region_area=0,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "# )\n",
    "# * SSA Light\n",
    "# mask_generator = SamAutomaticMaskGenerator(\n",
    "#     model=sam,\n",
    "#     points_per_side=16,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     crop_n_layers=0,  # 1 by default\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode='binary_mask',\n",
    "# )\n",
    "# * SSA Light 8点\n",
    "# mask_generator = SamAutomaticMaskGenerator(\n",
    "#     model=sam,\n",
    "#     points_per_side=8,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     crop_n_layers=0,  # 1 by default\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode='binary_mask',\n",
    "# )\n",
    "# * SSA Default\n",
    "# mask_generator = SamAutomaticMaskGenerator(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     crop_n_layers=0,  # 1 by default\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode='binary_mask',\n",
    "# )\n",
    "# * SSA Heavy\n",
    "# mask_generator = SamAutomaticMaskGenerator(\n",
    "#     model=sam,\n",
    "#     points_per_side=64,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode='binary_mask',\n",
    "# )\n",
    "# * SAM Auto阈值滤除分分、分。\n",
    "# mask_generator = SamAuto(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "#     score_thresh_offset=(-1.0, -1.0, 0.0),\n",
    "#     score_nms_offset=(0.0, 0.0, 0.0),\n",
    "#     stability_score_bias=(0.0, 0.0, 0.0),\n",
    "# )\n",
    "# * SAM Auto阈值滤除分分、分，并提高总级的阈值得分偏移。\n",
    "# mask_generator = SamAuto(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "#     score_thresh_offset=(-1.0, -1.0, 0.1),\n",
    "#     score_nms_offset=(0.0, 0.0, 0.0),\n",
    "#     stability_score_bias=(0.0, 0.0, 0.0),\n",
    "# )\n",
    "# * 降低总的稳定度阈值。\n",
    "# mask_generator = SamAuto(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "#     score_thresh_offset=(-1.0, -1.0, 0.1),\n",
    "#     score_nms_offset=(0.0, 0.0, 0.0),\n",
    "#     stability_score_bias=(0.0, 0.0, 0.04),\n",
    "# )\n",
    "# * SAM Auto总NMS优先。\n",
    "# mask_generator = SamAuto(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "#     score_thresh_offset=(0.0, 0.0, 0.0),\n",
    "#     score_nms_offset=(0.0, 0.0, 1.0),\n",
    "#     stability_score_bias=(0.0, 0.0, 0.0),\n",
    "# )\n",
    "# * SAM Auto总NMS优先，提高总级的阈值得分偏移。\n",
    "# mask_generator = SamAuto(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "#     score_thresh_offset=(0.0, 0.0, 0.1),\n",
    "#     score_nms_offset=(0.0, 0.0, 1.0),\n",
    "#     stability_score_bias=(0.0, 0.0, 0.0),\n",
    "# )\n",
    "# * SAM Auto总NMS优先，提高总级的阈值得分偏移，降低NMS阈值。\n",
    "# mask_generator = SamAuto(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.3,\n",
    "#     crop_n_layers=1,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "#     score_thresh_offset=(0.0, 0.0, 0.1),\n",
    "#     score_nms_offset=(0.0, 0.0, 1.0),\n",
    "#     stability_score_bias=(0.0, 0.0, 0.0),\n",
    "# )\n",
    "# * SAM Auto总NMS优先，提高总级的阈值得分偏移，crop0次。\n",
    "# mask_generator = SamAuto(\n",
    "#     model=sam,\n",
    "#     points_per_side=32,\n",
    "#     points_per_batch=256,\n",
    "#     pred_iou_thresh=0.86,\n",
    "#     stability_score_thresh=0.92,\n",
    "#     stability_score_offset=1.0,\n",
    "#     box_nms_thresh=0.7,\n",
    "#     crop_n_layers=0,\n",
    "#     crop_nms_thresh=0.7,\n",
    "#     crop_overlap_ratio=512 / 1500,\n",
    "#     crop_n_points_downscale_factor=2,\n",
    "#     min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "#     output_mode=\"binary_mask\",\n",
    "#     score_thresh_offset=(0.0, 0.0, 0.1),\n",
    "#     score_nms_offset=(0.0, 0.0, 1.0),\n",
    "#     stability_score_bias=(0.0, 0.0, 0.0),\n",
    "# )\n",
    "# * SAM Auto总NMS优先，提高总级的阈值得分偏移，结合石沙水过滤。\n",
    "mask_generator = SamAuto(\n",
    "    model=sam,\n",
    "    points_per_side=32,\n",
    "    points_per_batch=256,\n",
    "    pred_iou_thresh=0.86,\n",
    "    stability_score_thresh=0.92,\n",
    "    stability_score_offset=1.0,\n",
    "    box_nms_thresh=0.7,\n",
    "    crop_n_layers=1,\n",
    "    crop_nms_thresh=0.7,\n",
    "    crop_overlap_ratio=512 / 1500,\n",
    "    crop_n_points_downscale_factor=2,\n",
    "    min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    "    output_mode=\"binary_mask\",\n",
    "    score_thresh_offset=(0.0, 0.0, 0.1),\n",
    "    score_nms_offset=(0.0, 0.0, 1.0),\n",
    "    stability_score_bias=(0.0, 0.0, 0.0),\n",
    "    rock_sand_water_thresh=0.3,\n",
    "    rock_sand_water_chunk_size=50,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fb702ae3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-26T14:58:24.261372100Z",
     "start_time": "2023-05-26T14:57:29.646766800Z"
    }
   },
   "outputs": [],
   "source": [
    "img_masks = [mask_generator.generate(img) for img in imgs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c937160",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig: plt.Figure = plt.figure(dpi=200)\n",
    "show_imgs_anns(fig, imgs, img_ids, img_masks, levels=(0, 1, 2))\n",
    "fig.show()\n",
    "\n",
    "fig: plt.Figure = plt.figure(dpi=200)\n",
    "show_imgs_anns(fig, imgs, img_ids, img_masks, levels=(0, 1))\n",
    "fig.show()\n",
    "\n",
    "fig: plt.Figure = plt.figure(dpi=200)\n",
    "show_imgs_anns(fig, imgs, img_ids, img_masks, levels=(2,))\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
